import os
import time
import datetime
import pandas as pd
import platform
from multiprocessing import Pool, freeze_support, cpu_count
from joblib import Parallel, delayed
from tqdm import tqdm
import warnings

warnings.filterwarnings("ignore")
pd.set_option('expand_frame_repr', False)  # 当列太多时不换行
pd.set_option('display.max_rows', 8000)  # 最多显示数据的行数


def creat_stock_list(path): # 生成股票代码列表
    stock_list = []
    for root, dirs, files in os.walk(path): # os遍历文件夹中的所有文件
        if files:  # 当files不为空的时候
            for f in files:
                if f.endswith('.csv'):
                    stock_list.append(f[:8])
    return sorted(stock_list)


def update_stock(code):
    new_df = pd.read_csv(new_path + '/' + code + '.csv', skiprows = 1, parse_dates = ['交易日期'], encoding = 'gbk')
    path = old_path + '/' + code + '.csv'
    if os.path.exists(path): # 若有历史文件，则更新缺失部分
        old_df = pd.read_csv(path, skiprows = 1, parse_dates = ['交易日期'], encoding = "gbk")
        all_df = pd.concat([old_df, new_df])
        all_df.drop_duplicates(subset = ['交易日期'], keep = 'first', inplace = True)
        all_df.sort_values(by = ['交易日期'], inplace = True)
        if len(all_df) - len(old_df) != 0:
            pd.DataFrame(columns = [f'数据由{author}整理']).to_csv(path, index = False, encoding = 'gbk')
            all_df.to_csv(path, mode='a', encoding = 'gbk', index = False) # 保存的文件不含索引列
    else:
        if not new_df.empty:
            pd.DataFrame(columns = [f'数据由{author}整理']).to_csv(path, index = False, encoding = 'gbk')
            new_df.to_csv(path, mode='a', encoding = 'gbk', index = False) # 保存的文件不含索引列
    
# *************************************************** 调参数 ***************************************************
old_path = 'D:/python/xbx_stock_2021_part3/市值选谷策略深度优化实盘/data/daily_data/stock'
new_path = 'D:/python/xbx_stock_2021_part3/市值选谷策略深度优化实盘/data/daily_data/xbx增量日线' # 新数据文件夹
author = 'ziven'
# *************************************************** 主程序 ***************************************************
t1 = datetime.datetime.now()

old_list = creat_stock_list(old_path)
stock_list = creat_stock_list(new_path)
#stock_list = stock_list[0:25] # 用来测试
print('存量股票个数:', len(old_list))
print('增量股票个数:', len(stock_list))

if __name__ == '__main__': # 并行
    if 'Windows' in platform.platform(): # 添加对windows多进程的支持
        freeze_support()
    with Pool(max(cpu_count() - 1, 1)) as pool:
        Parallel(n_jobs=cpu_count() - 1)(delayed(update_stock)(code) for code in tqdm(stock_list)) # 适用spyder


t2 = datetime.datetime.now()
print(f"\n********** 全部完成 **********\n开始时间：{str(t1)[:-7]}\n结束时间：{str(t2)[:-7]}\n合计耗时：{str(t2-t1)[:-7]}")

